package me.helllp.dake.generator.engine.mybatis;

import com.baomidou.mybatisplus.annotation.IdType;
import com.baomidou.mybatisplus.annotation.TableField;
import com.baomidou.mybatisplus.annotation.TableId;
import com.squareup.javapoet.*;
import freemarker.template.TemplateException;
import lombok.Getter;
import lombok.Setter;
import me.helllp.dake.core.base.BasePo;
import me.helllp.dake.generator.engine.RdbmsGenerator;
import me.helllp.dake.generator.meta.ProjectMetaData;
import me.helllp.dake.generator.meta.core.SmartType;
import me.helllp.dake.generator.sql.api.IntrospectedTable;
import me.helllp.dake.generator.utils.FileAndPathUtils;
import me.helllp.dake.generator.utils.UnderlineToCamelUtils;

import javax.lang.model.element.Modifier;
import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;

/**
 * mybatis PO文件生成
 */
public class PoGenerator extends RdbmsGenerator {

    public PoGenerator(ProjectMetaData projectMetaData) {
        super(projectMetaData);
    }

    @Override
    public void generator() throws IOException, TemplateException {
        String packageName = getProjectMetaData().getBasePackage() + ".po";
        FileAndPathUtils.makePackage(getProjectMetaData().getSrcJavaPath(), packageName);

        getTables().stream().forEach(table -> {
            TypeSpec typeSpec = TypeSpec.classBuilder(UnderlineToCamelUtils.underlineToCamel(table.getTableName(), false))
                    .superclass(BasePo.class)
                    .addModifiers(Modifier.PUBLIC)
                    .addAnnotation(Getter.class)
                    .addAnnotation(Setter.class)
                    .addFields(parser(table))
                    .build();

            JavaFile file = JavaFile.builder(packageName, typeSpec).build();
            try {
                file.writeToFile(new File(getProjectMetaData().getSrcJavaPath()));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        });
    }

    private List<FieldSpec> parser(IntrospectedTable table) {
        List<FieldSpec> fieldSpecs = table.getPrimaryKeyColumns().stream().map(key -> {
            SmartType type = new SmartType(key.getFullyQualifiedJavaType().getPackageName(), key.getFullyQualifiedJavaType().getShortName());
            return FieldSpec.builder(type.build(),key.getJavaProperty(), Modifier.PRIVATE)
                    .addAnnotation(AnnotationSpec.builder(TableId.class)
                            .addMember("value", "$S", key.getActualColumnName())
                            .addMember("type", "$N.INPUT", ClassName.get(IdType.class).toString())
                            .build())
                    .build();
        }).collect(Collectors.toList());

        List<FieldSpec> tempSpecs = table.getBaseColumns().stream().map(key -> {
            SmartType type = new SmartType(key.getFullyQualifiedJavaType().getPackageName(), key.getFullyQualifiedJavaType().getShortName());
            return FieldSpec.builder(type.build(),key.getJavaProperty(), Modifier.PRIVATE)
                    .addAnnotation(AnnotationSpec.builder(TableField.class)
                            .addMember("value", "$S",key.getActualColumnName())
                            .build())
                    .build();
        }).collect(Collectors.toList());

        fieldSpecs.addAll(tempSpecs);

        return fieldSpecs;
    }
}
